# coding=utf-8
import os
import cv2
import imageio
import argparse
from PIL import Image


parser = argparse.ArgumentParser(description='img to vedio or gif')

parser.add_argument('--name', default='1_Plane_01',type=str, help='config file')
# parser.add_argument('--path', default='/home/xulong/snowstorm/0Dataset/VOT2018/godfather/godfather_result/',
parser.add_argument('--path', default='/home/xulong/snowstorm/0Dataset/rssrai2019/1_Plane_01/img/1_Plane_01/',
                    type=str, help='config file')

args = parser.parse_args()

# fps: 帧率
def make_video(video_path, fps):
    # fourcc = cv2.VideoWriter_fourcc(*"MJPG")
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")

    for image_name in os.listdir(root_path):
        img_path = os.path.join(root_path, image_name)
        break
    im = Image.open(img_path)
    print("im.size:", im.size)
    vw = cv2.VideoWriter(video_path, fourcc, fps, im.size)
    # vw = cv2.VideoWriter(path, 0x00000021, fps, im.size)
    filenames = os.listdir(root_path)
    filenames.sort(key=lambda x: int(x[:-4]))  #排序
    for image_name in filenames:
        frame = cv2.imread(root_path + '/' + image_name)
        vw.write(frame)
    print('Save video successfully!')

def make_gif(gif_path, duration):
    frames = []
    filenames = os.listdir(root_path)
    filenames.sort(key=lambda x: int(x[:-4]))  # 排序
    for image_name in filenames:
        if (len(image_name)) == 12:
            print('image_name:',image_name)
            frames.append(imageio.imread(root_path + '/' + image_name))
        else:
            continue
    imageio.mimsave(gif_path, frames, 'GIF', duration=duration)
    print('Save gif successfully!')
    return


if __name__ == '__main__':
    root_path = args.path
    video_name = args.name + '.mp4'
    gif_name = args.name + '.gif'
    # video_path = os.path.join(root_path, video_name)
    # gif_path = os.path.join(root_path, gif_name)
    video_path = os.path.join('/home/xulong/snowstorm/pysot_sgan_pytorch/results/remote_sensing/', video_name)
    gif_path = os.path.join('/home/xulong/snowstorm/pysot_sgan_pytorch/results/remote_sensing/', gif_name)
    # video_path = os.path.join('C:\\Users\\DZF\\Desktop\\anti_block_result\\elephant-2_f\\RPN++', video_name)
    # gif_path = os.path.join('C:\\Users\\DZF\\Desktop\\anti_block_result\\elephant-2_f\\RPN++', gif_name)

    vedio_fps = 15
    gif_duration = 1/vedio_fps

    make_video(video_path, vedio_fps)
    # make_gif(gif_path, gif_duration)


